{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "471686fe-a398-4268-8360-ed1b30f4884f",
   "metadata": {},
   "source": [
    "## OpenAI \n",
    "\n",
    "For this notebook we will use the OpenAI ChatGPT models. We import the OpenAI agent and set the api_key, you will have to provide your own API key. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "56f0ce8c-c65c-4b6b-932f-d88620d402de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup OpenAI Agent\n",
    "import openai\n",
    "\n",
    "openai.api_key = \"sk-your-key\"\n",
    "from llama_index.agent import OpenAIAgent"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9515f57c-408b-4ec2-97ba-048e5f89b7d2",
   "metadata": {},
   "source": [
    "## Database tool\n",
    "\n",
    "This tool connects to a database (using SQLAlchemy under the hood) and allows an Agent to query the database and get information about the tables.\n",
    "\n",
    "We import the ToolSpec and initialize it so that it can connect to our database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0cdc11ff-4cde-44ed-a2a3-5509d1582a11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import and initialize our tool spec\n",
    "from llama_hub.tools.database.base import DatabaseToolSpec\n",
    "\n",
    "db_spec = DatabaseToolSpec(\n",
    "    scheme=\"postgresql\",  # Database Scheme\n",
    "    host=\"localhost\",  # Database Host\n",
    "    port=\"5432\",  # Database Port\n",
    "    user=\"postgres\",  # Database User\n",
    "    password=\"x\",  # Database Password\n",
    "    dbname=\"dev_db\",  # Database Name\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "886a9be7-9048-43ac-86f0-7a69ed6ad408",
   "metadata": {},
   "source": [
    "After initializing the Tool Spec we have an instance of the ToolSpec. A ToolSpec can have many tools that it implements and makes available to agents. We can see the Tools by converting to the spec to a list of FunctionTools, using `to_tool_list`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3b370966-379a-4025-b201-855c115cc2e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load_data\n",
      "load_data(query: str) -> List[llama_index.schema.Document]\n",
      "Query and load data from the Database, returning a list of Documents.\n",
      "\n",
      "        Args:\n",
      "            query (str): an SQL query to filter tables and rows.\n",
      "\n",
      "        Returns:\n",
      "            List[Document]: A list of Document objects.\n",
      "        \n",
      "<class 'pydantic.main.load_data'>\n",
      "describe_tables\n",
      "describe_tables(tables: Optional[List[str]] = None) -> str\n",
      "\n",
      "            Describes the specifed tables in the database\n",
      "\n",
      "            Args:\n",
      "                tables (List[str]): A list of table names to retrieve details about\n",
      "        \n",
      "<class 'pydantic.main.describe_tables'>\n",
      "list_tables\n",
      "list_tables() -> List[str]\n",
      "\n",
      "            Returns a list of available tables in the database.\n",
      "            To retrieve details about the columns of specfic tables, use\n",
      "            the describe_tables endpoint\n",
      "        \n",
      "<class 'pydantic.main.list_tables'>\n"
     ]
    }
   ],
   "source": [
    "tools = db_spec.to_tool_list()\n",
    "for tool in tools:\n",
    "    print(tool.metadata.name)\n",
    "    print(tool.metadata.description)\n",
    "    print(tool.metadata.fn_schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "695a9655-d66e-4c4e-b660-b8a59e9ba5c8",
   "metadata": {},
   "source": [
    "We can see that the database tool spec provides 3 functions for the OpenAI agent. One to execute a SQL query, one to describe a list of tables in the database, and one to list all of the tables available in the database. \n",
    "\n",
    "We can pass the tool list to our OpenAI agent and test it out:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4e0a70d1-3160-42e2-ba56-75577f387175",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the Agent with our tools\n",
    "agent = OpenAIAgent.from_tools(tools, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "291184d2-ba13-471c-ad78-2b605bf06e6e",
   "metadata": {},
   "source": [
    "At this point our Agent is fully ready to start making queries to our database:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0ae36c93-ab77-4dd3-9cf0-a6ed7c3c598d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: list_tables with args: {}\n",
      "Got output: ['active_storage_blobs', 'addresses', 'amenities', 'amenities_dwellings', 'ar_internal_metadata', 'dwellings', 'listings', 'rooms', 'schema_migrations', 'users', 'waitlists', 'active_storage_attachments', 'active_storage_variant_records', 'messages', 'participants']\n",
      "========================\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Response(response='This database contains the following tables:\\n1. active_storage_blobs\\n2. addresses\\n3. amenities\\n4. amenities_dwellings\\n5. ar_internal_metadata\\n6. dwellings\\n7. listings\\n8. rooms\\n9. schema_migrations\\n10. users\\n11. waitlists\\n12. active_storage_attachments\\n13. active_storage_variant_records\\n14. messages\\n15. participants', source_nodes=[], metadata=None)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.chat(\"What tables does this database contain\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c77be8af-6e08-45da-b36f-8e90abb91cc6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: describe_tables with args: {\n",
      "  \"tables\": [\"messages\"]\n",
      "}\n",
      "Got output: \n",
      "CREATE TABLE messages (\n",
      "\tid BIGSERIAL NOT NULL, \n",
      "\tuser_id INTEGER NOT NULL, \n",
      "\troom_id INTEGER NOT NULL, \n",
      "\tbody TEXT, \n",
      "\tcreated_at TIMESTAMP(6) WITHOUT TIME ZONE NOT NULL, \n",
      "\tupdated_at TIMESTAMP(6) WITHOUT TIME ZONE NOT NULL, \n",
      "\tCONSTRAINT messages_pkey PRIMARY KEY (id), \n",
      "\tCONSTRAINT fk_rails_273a25a7a6 FOREIGN KEY(user_id) REFERENCES users (id), \n",
      "\tCONSTRAINT fk_rails_a8db0fb63a FOREIGN KEY(room_id) REFERENCES rooms (id)\n",
      ")\n",
      "\n",
      "\n",
      "\n",
      "========================\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Response(response='The `messages` table has the following columns:\\n\\n1. `id`: A unique identifier for each message (type: BIGSERIAL)\\n2. `user_id`: The ID of the user who sent the message (type: INTEGER)\\n3. `room_id`: The ID of the room where the message was sent (type: INTEGER)\\n4. `body`: The content of the message (type: TEXT)\\n5. `created_at`: The timestamp when the message was created (type: TIMESTAMP(6) WITHOUT TIME ZONE)\\n6. `updated_at`: The timestamp when the message was last updated (type: TIMESTAMP(6) WITHOUT TIME ZONE)\\n\\nThe table also has the following constraints:\\n- `messages_pkey`: Primary key constraint on the `id` column\\n- `fk_rails_273a25a7a6`: Foreign key constraint on the `user_id` column, referencing the `id` column of the `users` table\\n- `fk_rails_a8db0fb63a`: Foreign key constraint on the `room_id` column, referencing the `id` column of the `rooms` table', source_nodes=[], metadata=None)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.chat(\"Can you describe the messages table\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "54a57d0b-ebea-48e5-978c-90a7846b6504",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: load_data with args: {\n",
      "  \"query\": \"SELECT body FROM messages ORDER BY created_at DESC LIMIT 1\"\n",
      "}\n",
      "Got output: [Document(id_='10c3ba99-d12e-4557-805d-c7e95a9b1ee7', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, hash='80a5d3ba4f8f87aaac2096195bbd1d81f9ac58c75d85413da2d9fa533335789e', text='cool', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n')]\n",
      "========================\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Response(response='The most recent message has the following body: \"cool\".', source_nodes=[], metadata=None)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.chat(\"Fetch the most recent message and display the body\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ab1c135-3484-4d2e-8703-9cd5d5f8a1f0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
